package org.basis.algorithm.tree;

import org.basis.algorithm.tree.common.TreeNode;

import java.util.*;

/**
 * 求树的最大宽度
 *
 * @author Mr_wenpan@163.com 2021/12/20 20:46
 */
public class TreeMaxWidth {

    public static void main(String[] args) {

    }

    public static int maxWidth(TreeNode<Integer> head) {
        if (head == null) {
            return 0;
        }
        // 先求得每个节点属于第几层，记录在hash表
        HashMap<TreeNode<Integer>, Integer> map = new HashMap<>();
        LinkedList<TreeNode<Integer>> queue = new LinkedList<>();
        // 头节点压入队列，并设置在第0层
        queue.add(head);
        map.put(head, 0);

        while (!queue.isEmpty()) {
            TreeNode<Integer> pollFirst = queue.pollFirst();
            // 孩子节点入队列，并且将孩子节点的层数记录下来。孩子节点的所在层数为父节点的层数 + 1
            if (pollFirst.left != null) {
                queue.add(pollFirst.left);
                Integer parentLevel = map.get(pollFirst);
                map.put(pollFirst.left, parentLevel + 1);
            }
            if (pollFirst.right != null) {
                queue.add(pollFirst.right);
                Integer parentLevel = map.get(pollFirst);
                map.put(pollFirst.right, parentLevel + 1);
            }
        }

        // 在这里也可以直接通过map的相关api直接统计value出现次数最多的数字来实现最大宽度计算，但这不是算法练习的目的
//        List<Integer> list = new ArrayList<>( map.values());
//        Integer maxWidth = list.stream().max(Comparator.comparingInt(o -> o)).orElse(-1);

        // 再次宽度优先遍历，并统计每层的节点个数
        queue.add(head);
        // 当前所在层
        int currentLevel = 0;
        // 当前层的几点个数
        int currentLevelCount = 0;
        // 最大宽度
        int max = 0;

        while (!queue.isEmpty()) {
            // 弹出一个节点，并获取一下这个节点的所属层
            TreeNode<Integer> pollFirst = queue.pollFirst();
            Integer level = map.get(pollFirst);

            // 节点所属层 == 当前层，则当前层节点个数++
            if (level == currentLevel) {
                currentLevelCount++;
            } else {
                // 节点所属层 != 当前层，说明节点已经跨层了。重置当前层并将当前层节点个数值为1，并结算上一层的最大节点个数
                currentLevel++;
                max = Math.max(currentLevelCount, max);
                currentLevelCount = 1;
            }

            // 处理左右节点压入队列
            if (pollFirst.left != null) {
                queue.add(pollFirst.left);
            }
            if (pollFirst.right != null) {
                queue.add(pollFirst.right);
            }
        }

        return max;
    }
}
